from langchain_core.tools import tool
import pandas as pd
import numpy as np
from sklearn.metrics.pairwise import cosine_similarity
import json
from datetime import datetime

@tool
def similarity_calculation_tool(query: str) -> str:
    """
    计算司机相似度，基于历史接单行为特征。
    
    Args:
        query: JSON字符串，格式：{"input_file": "driver_data.csv", "target_driver_id": 1, "top_k": 10}
    
    Returns:
        相似司机列表和相似度分数
    """
    try:
        params = json.loads(query)
        input_file = params.get('input_file', 'driver_order_data.csv')
        target_driver_id = params.get('target_driver_id')
        top_k = params.get('top_k', 10)

        # 加载CSV数据
        file_path = f'ai_data/{input_file}'
        df = pd.read_csv(file_path)
        
        # 确保driver_id是整数类型
        df['driver_id'] = df['driver_id'].astype(int)

        if df.empty:
            return "相似度计算失败：没有找到司机数据"

        # 创建司机特征向量
        driver_features = df.groupby('driver_id').agg({
            'order_count': 'sum',
            'avg_amount': 'mean',
            'active_days': 'sum',
            'area_id': 'nunique',
            'time_category': lambda x: len(set(x))
        }).round(2)

        driver_features.columns = ['total_orders', 'avg_amount', 'active_days', 'unique_areas', 'time_categories']
        
        # 计算订单效率（总订单数/活跃天数）
        driver_features['order_efficiency'] = (driver_features['total_orders'] / driver_features['active_days']).fillna(0)
        
        # 计算区域效率（总订单数/唯一区域数）
        driver_features['area_efficiency'] = (driver_features['total_orders'] / driver_features['unique_areas']).fillna(0)

        # 标准化特征
        features_for_similarity = driver_features[['total_orders', 'avg_amount', 'active_days', 'unique_areas', 'time_categories', 'order_efficiency', 'area_efficiency']]
        
        # 处理无穷大和NaN值
        features_for_similarity = features_for_similarity.replace([np.inf, -np.inf], 0).fillna(0)
        
        # 计算相似度矩阵
        similarity_matrix = cosine_similarity(features_for_similarity)
        
        # 获取目标司机索引
        if target_driver_id is None:
            # 如果没有指定目标司机，选择订单数最多的司机
            target_driver_id = driver_features['total_orders'].idxmax()
        
        # 确保target_driver_id是整数
        target_driver_id = int(target_driver_id)
        
        if target_driver_id not in driver_features.index:
            # 如果目标司机不存在，选择订单数最多的司机作为替代
            print(f"目标司机ID {target_driver_id} 不存在，选择订单数最多的司机作为替代")
            target_driver_id = driver_features['total_orders'].idxmax()
            target_driver_id = int(target_driver_id)

        target_index = driver_features.index.get_loc(target_driver_id)
        
        # 获取相似度分数
        target_similarities = similarity_matrix[target_index]
        
        # 创建司机ID和相似度分数的对应关系
        driver_similarities = list(zip(driver_features.index, target_similarities))
        
        # 按相似度排序，排除自己
        driver_similarities = [(driver_id, score) for driver_id, score in driver_similarities if driver_id != target_driver_id]
        driver_similarities.sort(key=lambda x: x[1], reverse=True)
        
        # 获取前top_k个相似司机
        similar_drivers = []
        for driver_id, similarity_score in driver_similarities[:top_k]:
            driver_stats = driver_features.loc[driver_id]
            similar_drivers.append({
                'driver_id': int(driver_id),
                'similarity_score': float(similarity_score),
                'total_orders': int(driver_stats['total_orders']),
                'avg_amount': float(driver_stats['avg_amount']),
                'active_days': int(driver_stats['active_days']),
                'unique_areas': int(driver_stats['unique_areas']),
                'order_efficiency': float(driver_stats['order_efficiency'])
            })

        # 目标司机信息
        target_stats = driver_features.loc[target_driver_id]
        target_driver_info = {
            'total_orders': int(target_stats['total_orders']),
            'avg_amount': float(target_stats['avg_amount']),
            'active_days': int(target_stats['active_days']),
            'unique_areas': int(target_stats['unique_areas']),
            'order_efficiency': float(target_stats['order_efficiency'])
        }

        result = {
            'target_driver_id': int(target_driver_id),
            'target_driver_info': target_driver_info,
            'similar_drivers': similar_drivers,
            'feature_dimensions': len(features_for_similarity.columns),
            'analysis_date': datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        }

        return f"相似度计算完成：{json.dumps(result, ensure_ascii=False, indent=2)}"

    except Exception as e:
        return f"相似度计算失败：{str(e)}"
